{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "\n",
    "dataset_path = Path(\"datasets/raw\")\n",
    "\n",
    "def default_loader(path):\n",
    "    return Image.open(path).convert('RGB')\n",
    "\n",
    "class SkillDataset(Dataset):\n",
    "    def __init__(self, path: Path) -> None:\n",
    "        super().__init__()\n",
    "        all_img = list(dataset_path.glob(\"*.png\"));\n",
    "        self.c0 = [p for p in all_img if p.stem.endswith(\"_0\")]\n",
    "        self.c1 = [p for p in all_img if p.stem.endswith(\"_1\")]\n",
    "        self.c2 = [p for p in all_img if p.stem.endswith(\"_2\")]\n",
    "        self.c3 = [p for p in all_img if p.stem.endswith(\"_3\")]\n",
    "        print(\"len c0: \", len(self.c0), \" c1: \", len(self.c1), \" c2: \", len(self.c2), \" c3: \", len(self.c3))\n",
    "\n",
    "        self.loader = default_loader\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.GaussianBlur(3, sigma=(0.1, 2.0)),\n",
    "            transforms.RandomPosterize(3),\n",
    "            transforms.RandomAdjustSharpness(3),\n",
    "            transforms.RandomAutocontrast(),\n",
    "            transforms.ToTensor(),\n",
    "        ])\n",
    "        # 图片没多大，一次性全部载入内存算了\n",
    "        self.data = [ self.get(i) for i in range(len(self))]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.c0) + len(self.c1) + len(self.c2) + len(self.c3)\n",
    "    \n",
    "    def get(self, index):\n",
    "        if index % 1000 == 0:\n",
    "            print(\"load \", index, \" / \", len(self))\n",
    "        if index < len(self.c0):\n",
    "            return self.transform(self.loader(self.c0[index])), 0\n",
    "        elif index < len(self.c0) + len(self.c1):\n",
    "            return self.transform(self.loader(self.c1[index - len(self.c0)])), 1\n",
    "        elif index < len(self.c0) + len(self.c1) + len(self.c2):\n",
    "            return self.transform(self.loader(self.c2[index - len(self.c0) - len(self.c1)])), 2\n",
    "        else:\n",
    "            return self.transform(self.loader(self.c3[index - len(self.c0) - len(self.c1) - len(self.c2)])), 3\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        return self.data[index][0], self.data[index][1]\n",
    "        \n",
    "\n",
    "dataset = SkillDataset(dataset_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch\n",
    "\n",
    "train_size = int(0.8 * len(dataset))\n",
    "test_size = len(dataset) - train_size\n",
    "train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])\n",
    "\n",
    "# 显存不够可以把 batch size 改小点\n",
    "train_loader = DataLoader(train_dataset, batch_size=512, shuffle=True, num_workers=0)\n",
    "test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class ModelM5(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ModelM5, self).__init__()\n",
    "        self.channels = 3\n",
    "        self.classes = 4\n",
    "\n",
    "        self.conv1 = nn.Conv2d(self.channels, 32, 5, bias=False)\n",
    "        self.conv1_bn = nn.BatchNorm2d(32)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 5, bias=False)\n",
    "        self.conv2_bn = nn.BatchNorm2d(64)\n",
    "        self.conv3 = nn.Conv2d(64, 96, 5, bias=False)\n",
    "        self.conv3_bn = nn.BatchNorm2d(96)\n",
    "        self.conv4 = nn.Conv2d(96, 128, 5, bias=False)\n",
    "        self.conv4_bn = nn.BatchNorm2d(128)\n",
    "        self.conv5 = nn.Conv2d(128, 160, 5, bias=False)\n",
    "        self.conv5_bn = nn.BatchNorm2d(160)\n",
    "        self.fc1 = nn.Linear(924160, self.classes, bias=False)\n",
    "        self.fc1_bn = nn.BatchNorm1d(self.classes)\n",
    "    def get_logits(self, x):\n",
    "        x = (x - 0.5) * 2.0\n",
    "        conv1 = F.relu(self.conv1_bn(self.conv1(x)))\n",
    "        conv2 = F.relu(self.conv2_bn(self.conv2(conv1)))\n",
    "        conv3 = F.relu(self.conv3_bn(self.conv3(conv2)))\n",
    "        conv4 = F.relu(self.conv4_bn(self.conv4(conv3)))\n",
    "        conv5 = F.relu(self.conv5_bn(self.conv5(conv4)))\n",
    "        flat5 = torch.flatten(conv5.permute(0, 2, 3, 1), 1)\n",
    "        logits = self.fc1_bn(self.fc1(flat5))\n",
    "        return logits\n",
    "    def forward(self, x):\n",
    "        logits = self.get_logits(x)\n",
    "        return F.log_softmax(logits, dim=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from torch.cuda.amp import autocast, GradScaler\n",
    "\n",
    "use_cuda = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "if not use_cuda:\n",
    "    print(\"WARNING: CPU will be used for training.\")\n",
    "\n",
    "model = ModelM5().to(device)\n",
    "\n",
    "criterion = torch.nn.CrossEntropyLoss().to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "scaler = GradScaler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "start_time = time.time()\n",
    "def train(epoch):\n",
    "    global start_time\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        with autocast():\n",
    "            output = model(data)\n",
    "            loss = criterion(output, target)\n",
    "            \n",
    "        scaler.scale(loss).backward()\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "        \n",
    "    cur_time = time.time()\n",
    "    cost = cur_time - start_time\n",
    "    print(f'Train Epoch: {epoch}, Loss: {loss.item():.8f}, cost: {cost:.2f} s')\n",
    "    start_time = cur_time\n",
    "            \n",
    "def test():\n",
    "    model.eval()\n",
    "    test_loss = 0.0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            with autocast():\n",
    "                output = model(data)\n",
    "                test_loss += criterion(output, target).item()\n",
    "            pred = output.argmax(dim=1, keepdim=True)\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    acc = 100. * correct / len(test_loader.dataset)\n",
    "\n",
    "    print(f'=== Test: Loss: {test_loss:.8f}, Acc: {acc:.4f} ===')\n",
    "    return test_loss, acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_cuda:\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "output = Path('checkpoints')\n",
    "output.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "model_name = model.__class__.__name__\n",
    "best_model_path = output / f'{model_name}_best.pt'\n",
    "\n",
    "best_epoch = 0\n",
    "best_loss = 100.0\n",
    "best_acc = 0.0\n",
    "default_interval = 1\n",
    "\n",
    "def pipeline(start_epoch = 0, test_interval = default_interval):\n",
    "    global best_epoch, best_loss, best_acc\n",
    "\n",
    "    for epoch in range(start_epoch, 1000):\n",
    "        train(epoch)\n",
    "        if epoch % test_interval != 0:\n",
    "            continue\n",
    "        \n",
    "        loss, acc = test()\n",
    "        print(f'=== Pre best is {best_epoch}, Loss: {best_loss:.8f}, Acc: {best_acc:.4f} ===')\n",
    "        torch.save(model, output / f'{model_name}_{epoch}.pt')\n",
    "        if loss > best_loss:\n",
    "            if epoch - best_epoch > 100:\n",
    "                print('No improvement for a long time, Early stop!')\n",
    "                break\n",
    "            else:\n",
    "                continue\n",
    "        best_epoch = epoch\n",
    "        best_loss = loss\n",
    "        best_acc = acc\n",
    "        print(f'====== New best is {best_epoch}, Loss: {best_loss:.8f}, Acc: {best_acc:.4f} ======')\n",
    "        torch.save(model, best_model_path)\n",
    "\n",
    "pipeline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.load(best_model_path)\n",
    "test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导出 onnx\n",
    "\n",
    "import torch.onnx\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "def convert_onnx(path: Path):\n",
    "    model = torch.load(path, map_location=torch.device(\"cpu\"))\n",
    "    model.eval()\n",
    "    dummy_input = torch.randn(1, 3, 96, 96)\n",
    "    torch.onnx.export(\n",
    "        model,\n",
    "        dummy_input,\n",
    "        path.with_suffix(\".onnx\"),\n",
    "        input_names=[\"input\"],\n",
    "        output_names=[\"output\"],\n",
    "    )\n",
    "\n",
    "\n",
    "convert_onnx(best_model_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "\n",
    "onnx.checker.check_model(str(best_model_path.with_suffix(\".onnx\")))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chat",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
